# koroFileHeader at Yungoal acer
# Create: 2021-04-11 10:21:58
# LastEdit: 2021-09-30 13:36:52
"""请求操作使用的 client 实例(单例)"""
__author__ = '749B'

from datetime import timedelta
from dateutil.parser import parse as date_parse
import json
import os
import requests
import time
from urllib.parse import urljoin

import azure.mgmt.resourcegraph as arg
from azure.identity import DefaultAzureCredential
from azure.mgmt.resource import SubscriptionClient
from azure.mgmt.resourcegraph import ResourceGraphClient

from utils.tools import init_printer

from typing import List, Mapping, Dict, Any, Optional, NoReturn, Tuple, Union
# 只是因为函数的返回需要做类型判断
from azure.core.pipeline.transport import HttpRequest

class Azure_Client_Base(object):
    """
    Client 基类：都设置为单例模式
    """

    _instance = None
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, cloud, credential):
        self._client = None
        self.cloud = cloud
        self.credential = credential


class Azure_Client(Azure_Client_Base):

    def __init__(self, cloud, /, **kwargs):
        self.cloud = cloud
        self.credential = DefaultAzureCredential(
            authority=self.cloud.endpoints.active_directory, 
            **kwargs)
        
        self.subscription_client = Azure_Subscription_Client(self.cloud, self.credential)
        self.add_func(self.subscription_client)
        self.resource_graph_client = Azure_Resource_Graph_Client(self.cloud, self.credential)
        self.add_func(self.resource_graph_client)
        self.restapi_client = Azure_Restapi_Client(self.cloud)
        self.add_func(self.restapi_client)

    def add_func(self, obj: Azure_Client_Base) -> None:
        """将实例obj的方法变成自己实例的方法"""
        for name in dir(obj):
            if name.startswith('_'):
                continue
            if name in ['client', 'token']:
                # 跳过 @property 设置的属性方法
                # 否则getattr的时候会去执行取到值，这样就没有延迟加载的意义了
                continue
            if hasattr(self, name):
                continue
            func = getattr(obj, name)
            if not callable(func):
                continue
            setattr(self, name, func)


class Azure_Subscription_Client(Azure_Client_Base):
    
    @property
    def client(self) -> SubscriptionClient:
        """延迟实例化client"""
        if self._client:
            return self._client

        self._client = SubscriptionClient(
            self.credential,
            base_url=self.cloud.endpoints.resource_manager,
            credential_scopes=[self.cloud.endpoints.resource_manager + "/.default"])
        return self._client

    def get_subscriptions(self) -> List[Dict[str, Any]]:
        """如果只要提取id的列表
        [sub.get("subscription_id") for sub in client.get_subscriptions()]
        """
        return [sub.as_dict() for sub in self.client.subscriptions.list()]

    def get_subscriptionid_by_name(self, name: str, accept_id: bool = True) -> Optional[str]:
        """提供订阅的名称，返回订阅的id
        accept_id: 默认开启，如果直接给id，并且有这个id，就直接返回id
        """
        map = {sub.get("display_name"): sub.get("subscription_id") for sub in self.get_subscriptions()}
        if accept_id and name in map.values():
            return name
        return map.get(name)

    def subscription_display_name(self) -> str:
        """返回字符串，可以把订阅ID转换成订阅的名称
        返回的格式：
        | extend subscriptionDisplayName=case(subscriptionId =~ 'ID1','NAME1',subscriptionId =~ 'ID2','NAME2',subscriptionId)

        把返回的字符串，添加到查询语句里，单独做一行。
        然后 `| project subscriptionDisplayName` 来显示这个字段。
        用法示例：
        subscriptionDisplayName = client.subscription_display_name()
        strQuery = textwrap.dedent('''\\
        recoveryservicesresources
        | where type =~ "microsoft.recoveryservices/vaults/backupfabrics/protectioncontainers/protecteditems"
        {subscriptionDisplayName}
        | project subscriptionDisplayName, resourceGroup, properties.lastRecoveryPoint, properties.lastBackupTime
        ''').format(
            subscriptionDisplayName=subscriptionDisplayName, 
            )
        """
        str_list = []
        for subs in self.get_subscriptions():
            s1 = f"subscriptionId =~ '{subs.get('subscription_id')}','{subs.get('display_name')}'" 
            str_list.append(s1)
        subscriptionDisplayName = "| extend subscriptionDisplayName=case({},subscriptionId)".format(
            ','.join(str_list)
        )
        return subscriptionDisplayName


class Azure_Resource_Graph_Client(Azure_Client_Base):

    @property
    def client(self) -> ResourceGraphClient:
        """延迟实例化client"""
        if self._client:
            return self._client
            
        self._client = ResourceGraphClient(
            self.credential,
            base_url=self.cloud.endpoints.resource_manager,
            credential_scopes=[self.cloud.endpoints.resource_manager + "/.default"])
            
        # 修改掉源码中的一个方法
        def my_post(
            url,  # type: str
            params=None,  # type: Optional[Dict[str, str]]
            headers=None,  # type: Optional[Dict[str, str]]
            content=None,  # type: Any
            form_content=None,  # type: Optional[Dict[str, Any]]
            stream_content=None,  # type: Any
        ):
            # type: (...) -> HttpRequest
            if params.get('api-version', None):
                params['api-version'] = '2019-04-01'
            return _old_post(url, params, headers, content, form_content, stream_content)

        _old_post = self._client._client.post
        self._client._client.post = my_post
        # 修改方法结束

        return self._client

    def resource_graph_query(
        self, 
        options_kwargs: Mapping[str, Any], 
        request_kwargs: Mapping[str, Any], 
        resources_kwargs: Mapping[str, Any]
        ) -> arg.models.QueryResponse:
        """
        大多数参数都不用定制，使用resource_graph_query_default仅需要设置基本的参数
        """
        argQueryOptions = arg.models.QueryRequestOptions(**options_kwargs) 
        argQuery = arg.models.QueryRequest(options=argQueryOptions, **request_kwargs)
        argResults = self.client.resources(argQuery, **resources_kwargs)
        return argResults

    def resource_graph_query_default(
        self, 
        strQuery: str, 
        subscriptions: List[str], 
        result_format: str = "objectArray"
        ) -> arg.models.QueryResponse:
        """
        result_format: objectArray or table. 也就objectArray合适
        """
        return self.resource_graph_query(
            {"result_format": result_format},  # QueryRequestOptions
            {"subscriptions": subscriptions, "query": strQuery}, # QueryRequest
            {},  # resources
        )


class Azure_Restapi_Client(Azure_Client_Base):
    
    def __init__(self, cloud):
        self.cloud = cloud

        self._token = None
    
    @property
    def token(self) -> Union[str, NoReturn]:
        """延迟获取token"""
        if self._token:
            return self._token

        base_url=self.cloud.endpoints.active_directory
        url_path = '{azure_tenant_id}/oauth2/v2.0/token'.format(
            azure_tenant_id=os.environ.get('AZURE_TENANT_ID')
        )
        url = urljoin(base_url, url_path)
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        data = {
            'client_id': os.environ.get('AZURE_CLIENT_ID'),
            'scope': self.cloud.endpoints.resource_manager + "/.default",
            'client_secret': os.environ.get('AZURE_CLIENT_SECRET'),
            'grant_type': 'client_credentials',
        }
        r = requests.post(url, data=data, headers=headers)
        r.raise_for_status()
        self._token = r.json().get('access_token')
        return self._token

    def runcmd_run(
        self, 
        subscription_id: str, 
        resource_group_name: str, 
        vm_name: str, 
        script: List[str],
        api_version: str = '2021-03-01',
        parameters: Optional[List[Mapping[str, str]]] = None,
        is_linux: bool = True
        ) -> requests.Response:
        """执行命令"""
        base_url = self.cloud.endpoints.resource_manager
        url_path = f'subscriptions/{subscription_id}/' \
            f'resourceGroups/{resource_group_name}/' \
            f'providers/Microsoft.Compute/virtualMachines/{vm_name}/' \
            f'runCommand?api-version={api_version}'
        url = urljoin(base_url, url_path)
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        command_id = 'RunShellScript' if is_linux else 'RunPowerShellScript'
        data = {
            'commandId': command_id,
            'script': script,
            }
        if parameters:
            data['parameters'] = parameters

        r = requests.post(url, data=json.dumps(data), headers=headers)
        # r.status_code == 202: 
        #   期待的结果，返回获取结果的地址
        #   r.headers.get('Azure-AsyncOperation')
        # r.status_code == 409:
        #   错误 message 如下，上一个命令还没执行完毕
        #   Run command extension execution is in progress. 
        #   Please wait for completion before invoking a run command.
        return r

    def runcmd_result(self, async_operation: str) -> requests.Response:
        """异步拿结果"""
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        return requests.get(async_operation, headers=headers)

    def runcmd_onestep(self, *args, report=False, wait=15, timeout=300, **kwargs) -> Tuple[bool, requests.Response]:
        """一步获取结果，但是不好用
        report: 调试的时候，打开。避免程序长时间没有响应
        wait: 根据脚本测试的情况，设置一个差不多的值。默认先等待一段时间
        需要在主程序中用多线程来执行，但是每个线程执行时的异常是会被多线程模块捕获的。
        调试起来不方便，效率也差。
        """
        # 并发安全的打印
        printer = init_printer()
        r = self.runcmd_run(*args, **kwargs)
        if r.text:
            return False, r
        assert r.status_code == 202
        async_operation = r.headers.get('Azure-AsyncOperation')
        
        # 运行脚本的最短时间大约为 20 秒。
        vm_name = 'Not Found'
        if report:
            if len(args) >= 3:
                vm_name = args[2]
            else:
                vm_name = kwargs.get('vm_name')
            printer(f"{vm_name}: 准备sleep {wait} 秒")
        time.sleep(wait)
        retry = 1
        res = self.runcmd_result(async_operation)
        while res.ok and retry < timeout:
            status = res.json().get('status')
            if status.lower() == 'succeeded':
                break
            if report:
                printer(f"{vm_name}: 准备sleep {retry} 秒")
            time.sleep(retry)
            retry <<= 1
            res = self.runcmd_result(async_operation)
        status = res.json().get('status')
        if status.lower() == 'succeeded':
            return True, res
        return False, res

    def resourcehealth_availabilitystatuses(self):
        """
        资源当前使用状态
        https://docs.microsoft.com/zh-cn/rest/api/resourcehealth/availabilitystatuses/getbyresource
        虚拟机开机状态: {
        'availabilityState': 'Available',
        'title': 'Available',
        'summary': "There aren't any known Azure platform problems "
                    'affecting this virtual machine.'}
        虚拟机关机状态: {
        'availabilityState': 'Unavailable',
        'title': 'Stopped and deallocated',
        'summary': 'This virtual machine is stopped and deallocated as '
                    'requested by an authorized user or process.'}
        """
        pass

    def backupjobs_list(
        self,
        subscription_id: str,
        resource_group_name: str,
        vault_name: str,
        api_version: str = '2021-02-10',
        filter: Optional[str] = None,
        skip_token: Optional[str] = None,
        next_url: Optional[str] = None,
        ) -> requests.Response:
        """恢复服务 - 备份/备份作业/列表
        https://docs.microsoft.com/zh-cn/rest/api/backup/2021-02-10/backup-jobs/list
        skip_token: 文档有这个参数，但是没例子，具体作用不清楚。也可能就是给nextLink用的，而不是让用户直接使用的
        如果时间跨度太长，就会导致返回结果太多，无法一次全部返回。返回结果会附加一个nextLink。这时，直接给nextLink发Get请求即可
        """
        if next_url:
            url = next_url
        else:
            base_url = self.cloud.endpoints.resource_manager
            url_path = f'subscriptions/{subscription_id}/' \
                f'resourceGroups/{resource_group_name}/' \
                'providers/Microsoft.RecoveryServices/' \
                f'vaults/{vault_name}/' \
                f'backupJobs?api-version={api_version}'
            if filter:
                url_path += f'&$filter={filter}'
            if skip_token:
                url_path += f'&$skipToken={skip_token}'
            url = urljoin(base_url, url_path)
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        r = requests.get(url, headers=headers)
        return r

    @staticmethod
    def backupjobs_filter(
        status: Optional[str] = 'Completed',
        backupManagementType: Optional[str] = 'AzureIaasVM',
        operation: Optional[str] = 'Backup',
        jobId: Optional[str] = None,
        startTime: Optional[str] = None,
        endTime: Optional[str] = None,
        # 下面是非基本属性
        start_day: Optional[str] = None,
        end_day: Optional[str] = None,
        ) -> Optional[str]:
        """
        支持的具体过滤条件和内容如下：
        status: InProgress, Completed, Failed, CompletedWithWarnings, Cancelled, Cancelling
        backupManagementType: AzureIaasVM, MAB, DPM, AzureBackupServer, AzureSql
        operation: ConfigureBackup, Backup, Restore, DisableBackup, DeleteBackupData
        jobId: guid
        startTime: yyyy-mm-dd hh:mm:ss PM
        endTime: yyyy-mm-dd hh:mm:ss PM
        """
        filters = []
        if status:
            filters.append(f"status eq '{status}'")
        if backupManagementType:
            filters.append(f"backupManagementType eq '{backupManagementType}'")
        if operation:
            filters.append(f"operation eq '{operation}'")
        if jobId:
            filters.append(f"jobId eq '{jobId}'")
        if startTime:
            filters.append(f"startTime eq '{startTime}'")
        elif start_day:
            start_day_obj = date_parse(start_day)
            filters.append(f"startTime eq '{start_day_obj.year}-{start_day_obj.month:02}-{start_day_obj.day:02} 00:00:00 AM'")
        if endTime:
            filters.append(f"endTime eq '{endTime}'")
        elif end_day:
            end_day_obj = date_parse(end_day) + timedelta(days=1)
            filters.append(f"endTime eq '{end_day_obj.year}-{end_day_obj.month:02}-{end_day_obj.day:02} 00:00:00 AM'")
        if filter:
            return ' and '.join(filters)
        return None

    def backupjobs_yield_all(
        self,
        subscription_id: str,
        resource_group_name: str,
        vault_name: str,
        api_version: str = '2021-02-10',
        start: Optional[str] = None,
        end: Optional[str] = None,
        ) -> requests.Response:
        """遍历所有的记录
        API调用是，不提供开始/结束时间，就是全部记录了
        """
        base_url = self.cloud.endpoints.resource_manager
        url_path = f'subscriptions/{subscription_id}/' \
            f'resourceGroups/{resource_group_name}/' \
            'providers/Microsoft.RecoveryServices/' \
            f'vaults/{vault_name}/' \
            f'backupJobs?api-version={api_version}'
        filter = self.backupjobs_filter(start_day=start, end_day=end)
        if filter:
            url_path += f'&$filter={filter}'
        url = urljoin(base_url, url_path)
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        r = requests.get(url, headers=headers)
        if not r:
            print(r.json())
        while r:
            yield r
            next_link = r.json().get('nextLink')
            if not next_link:
                break
            r = requests.get(next_link, headers=headers)
        return

    def monitor_metric_definitions(
        self,
        subscription_id: str,
        resource_group_name: str,
        resource_type: str,
        resource_name: str,
        api_version: str = '2018-01-01',
        params: Optional[Mapping[str, str]] = None,
        ) -> requests.Response:
        """检索指标定义
        https://docs.microsoft.com/zh-cn/rest/api/monitor/metric-definitions/list
        """
        base_url = self.cloud.endpoints.resource_manager
        url_path = f'subscriptions/{subscription_id}/' \
            f'resourceGroups/{resource_group_name}/' \
            f'providers/{resource_type}/{resource_name}/' \
            f'providers/microsoft.insights/' \
            f'metricDefinitions?api-version={api_version}'
        url = urljoin(base_url, url_path)
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        r = requests.get(url, headers=headers, params=params)
        return r

    def monitor_metric(
        self,
        subscription_id: str,
        resource_group_name: str,
        resource_type: str,
        resource_name: str,
        api_version: str = '2018-01-01',
        params: Optional[Mapping[str, str]] = None,
        ) -> requests.Response:
        """检索指标值
        https://docs.microsoft.com/zh-cn/rest/api/monitor/metrics/list
        直接用资源的id也可以，这样就不用拼接各种资源信息了。用下面提供的另外的方法
        """
        base_url = self.cloud.endpoints.resource_manager
        url_path = f'subscriptions/{subscription_id}/' \
            f'resourceGroups/{resource_group_name}/' \
            f'providers/{resource_type}/{resource_name}/' \
            f'providers/microsoft.insights/' \
            f'metrics?api-version={api_version}'
        url = urljoin(base_url, url_path)
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        r = requests.get(url, headers=headers, params=params)
        return r

    def monitor_metric_resourceid(
        self,
        resource_id: str,
        api_version: str = '2018-01-01',
        params: Optional[Mapping[str, str]] = None,
        ) -> requests.Response:
        """检索指标值
        https://docs.microsoft.com/zh-cn/rest/api/monitor/metrics/list
        提供另外一种直接提供资源id的调用方式，这个更直接更方便
        """
        base_url = self.cloud.endpoints.resource_manager
        url_path = f'{resource_id}/' \
            f'providers/microsoft.insights/' \
            f'metrics?api-version={api_version}'
        url = urljoin(base_url, url_path)
        headers = {
            'Authorization': 'Bearer {token}'.format(token=self.token),
            'Content-Type': 'application/json',
        }
        r = requests.get(url, headers=headers, params=params)
        return r

